import numpy as np
import pandas as pd
import re

# CYP酶频率数据
CYP_FREQUENCIES = {
    'CYP1A1': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP1A2': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP2A6': {'EM': 0.537, 'PM': 0.015, 'IM1': 0.399, 'IM2': 0.049, 
               'UM': 0.0},
    'CYP2B6': {'EM': 0.524, 'PM': 0.068, 'IM1': 0.323, 'IM2': 0.0, 
               'UM': 0.085},
    'CYP2C8': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP2C9': {'EM': 0.934, 'PM': 0.003, 'IM1': 0.063, 'IM2': 0.0, 
               'UM': 0.0},
    'CYP2C18': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP2C19': {'EM': 0.396, 'PM': 0.134, 'IM1': 0.458, 'IM2': 0.012, 
                'UM': 0.0},
    'CYP2D6': {'EM': 0.597, 'PM': 0.003, 'IM1': 0.39, 'IM2': 0.0, 
               'UM': 0.01},
    'CYP2E1': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP2J2': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP3A4': {'EM': 1.0, 'PM': 0.0, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP3A5': {'EM': 0.42, 'PM': 0.58, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0},
    'CYP3A7': {'EM': 0.12, 'PM': 0.88, 'IM1': 0.0, 'IM2': 0.0, 'UM': 0.0}
}

# CYP酶丰度和周转数据
CYP_ABUNDANCES = {
    'CYP1A1': {
        'EM': {'mean': 1.24, 'cv': 118}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0183, 'turnover_cv': 56
    },
    'CYP1A2': {
        'EM': {'mean': 42, 'cv': 50}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0183, 'turnover_cv': 56
    },
    'CYP2A6': {
        'EM': {'mean': 18.8, 'cv': 57}, 'PM': {'mean': 1.8, 'cv': 57},
        'IM1': {'mean': 13, 'cv': 57}, 'IM2': {'mean': 7.7, 'cv': 57},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0267, 'turnover_cv': 56
    },
    'CYP2B6': {
        'EM': {'mean': 6.7, 'cv': 63}, 'PM': {'mean': 1.4, 'cv': 85},
        'IM1': {'mean': 5, 'cv': 105}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 7.7, 'cv': 128}, 'turnover_mean': 0.0217, 'turnover_cv': 56
    },
    'CYP2C8': {
        'EM': {'mean': 7.7, 'cv': 75}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0301, 'turnover_cv': 56
    },
    'CYP2C9': {
        'EM': {'mean': 87.6, 'cv': 55}, 'PM': {'mean': 40.5, 'cv': 88},
        'IM1': {'mean': 76.7, 'cv': 81}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0067, 'turnover_cv': 56
    },
    'CYP2C18': {
        'EM': {'mean': 0.4, 'cv': 39}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0267, 'turnover_cv': 56
    },
    'CYP2C19': {
        'EM': {'mean': 4.4, 'cv': 52}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 2.85, 'cv': 52}, 'IM2': {'mean': 7.01, 'cv': 89},
        'UM': {'mean': 10.23, 'cv': 79}, 'turnover_mean': 0.0267, 'turnover_cv': 56
    },
    'CYP2D6': {
        'EM': {'mean': 10.47, 'cv': 65}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 3.29, 'cv': 65}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 20.25, 'cv': 65}, 'turnover_mean': 0.0099, 'turnover_cv': 56
    },
    'CYP2E1': {
        'EM': {'mean': 124, 'cv': 68}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0176, 'turnover_cv': 63
    },
    'CYP2J2': {
        'EM': {'mean': 2, 'cv': 75}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0194, 'turnover_cv': 56
    },
    'CYP3A4': {
        'EM': {'mean': 120, 'cv': 33}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0193, 'turnover_cv': 68
    },
    'CYP3A5': {
        'EM': {'mean': 82.3, 'cv': 68}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0193, 'turnover_cv': 68
    },
    'CYP3A7': {
        'EM': {'mean': 14, 'cv': 71}, 'PM': {'mean': 0, 'cv': 0},
        'IM1': {'mean': 0, 'cv': 0}, 'IM2': {'mean': 0, 'cv': 0},
        'UM': {'mean': 0, 'cv': 0}, 'turnover_mean': 0.0193, 'turnover_cv': 68
    }
}

# 转运体频率数据
TRANSPORTER_FREQUENCIES = {
    'ABCB1 (P-gp/MDR1)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'ABCB11 (BSEP)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'ABCC2 (MRP2)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'ABCC3 (MRP3)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'ABCC4 (MRP4)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'ABCC6 (MRP6)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'ABCG2 (BCRP)': {'ET': 0.44, 'PT': 0.1, 'IT': 0.46, 'UT': 0.0},
    'SLC29A1 (ENT1)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC29A2 (ENT2)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC10A1 (NTCP)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC16A1 (MCT1)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLCO1B1 (OATP1B1)': {'ET': 0.584, 'PT': 0.21, 'IT': 0.4, 'UT': 0.031},
    'SLCO1B3 (OATP1B3)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLCO2B1 (OATP2B1)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC22A1 (OCT1)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC22A7 (OAT2)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC22A9 (OAT7)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC47A1 (MATE1)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0},
    'SLC51A/B (OST-α/β)': {'ET': 1.0, 'PT': 0.0, 'IT': 0.0, 'UT': 0.0}
}

# 转运体丰度和周转数据（第一列是绝对丰度，其他列是相对于第一列的比例）
TRANSPORTER_ABUNDANCES = {
    'ABCB1 (P-gp/MDR1)': {
        'ET': {'mean': 0.246, 'cv': 59}, 'PT': {'mean': 0.246, 'cv': 59},
        'IT': {'mean': 0.246, 'cv': 59}, 'UT': {'mean': 0.246, 'cv': 59},
        'turnover_mean': 0.054, 'turnover_cv': 28
    },
    'ABCB11 (BSEP)': {
        'ET': {'mean': 1.0, 'cv': 131}, 'PT': {'mean': 1.0, 'cv': 131},
        'IT': {'mean': 1.0, 'cv': 131}, 'UT': {'mean': 1.0, 'cv': 131},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'ABCC2 (MRP2)': {
        'ET': {'mean': 0.59, 'cv': 88}, 'PT': {'mean': 0.59, 'cv': 88},
        'IT': {'mean': 0.59, 'cv': 88}, 'UT': {'mean': 0.59, 'cv': 88},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'ABCC3 (MRP3)': {
        'ET': {'mean': 0.239, 'cv': 65}, 'PT': {'mean': 0.239, 'cv': 65},
        'IT': {'mean': 0.239, 'cv': 65}, 'UT': {'mean': 0.239, 'cv': 65},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'ABCC4 (MRP4)': {
        'ET': {'mean': 0.0, 'cv': 60}, 'PT': {'mean': 0.0, 'cv': 60},
        'IT': {'mean': 0.0, 'cv': 60}, 'UT': {'mean': 0.0, 'cv': 60},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'ABCC6 (MRP6)': {
        'ET': {'mean': 0.214, 'cv': 46}, 'PT': {'mean': 0.214, 'cv': 46},
        'IT': {'mean': 0.214, 'cv': 46}, 'UT': {'mean': 0.214, 'cv': 46},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'ABCG2 (BCRP)': {
        'ET': {'mean': 0.103, 'cv': 30}, 'PT': {'mean': 0.37, 'cv': 30},
        'IT': {'mean': 0.67, 'cv': 30}, 'UT': {'mean': 0.0, 'cv': 30},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC29A1 (ENT1)': {
        'ET': {'mean': 0.0646, 'cv': 49}, 'PT': {'mean': 0.0646, 'cv': 49},
        'IT': {'mean': 0.0646, 'cv': 49}, 'UT': {'mean': 0.0646, 'cv': 49},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC29A2 (ENT2)': {
        'ET': {'mean': 0.0, 'cv': 60}, 'PT': {'mean': 0.0, 'cv': 60},
        'IT': {'mean': 0.0, 'cv': 60}, 'UT': {'mean': 0.0, 'cv': 60},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC10A1 (NTCP)': {
        'ET': {'mean': 0.647, 'cv': 50}, 'PT': {'mean': 0.647, 'cv': 50},
        'IT': {'mean': 0.647, 'cv': 50}, 'UT': {'mean': 0.647, 'cv': 50},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC16A1 (MCT1)': {
        'ET': {'mean': 0.638, 'cv': 60}, 'PT': {'mean': 0.638, 'cv': 60},
        'IT': {'mean': 0.638, 'cv': 60}, 'UT': {'mean': 0.638, 'cv': 60},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLCO1B1 (OATP1B1)': {
        'ET': {'mean': 3.1, 'cv': 73}, 'PT': {'mean': 0.21, 'cv': 29},
        'IT': {'mean': 0.4, 'cv': 51}, 'UT': {'mean': 0.81, 'cv': 62},
        'turnover_mean': 0.031, 'turnover_cv': 27
    },
    'SLCO1B3 (OATP1B3)': {
        'ET': {'mean': 3.08, 'cv': 89}, 'PT': {'mean': 3.08, 'cv': 89},
        'IT': {'mean': 3.08, 'cv': 89}, 'UT': {'mean': 3.08, 'cv': 89},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLCO2B1 (OATP2B1)': {
        'ET': {'mean': 1.18, 'cv': 62}, 'PT': {'mean': 1.18, 'cv': 62},
        'IT': {'mean': 1.18, 'cv': 62}, 'UT': {'mean': 1.18, 'cv': 62},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC22A1 (OCT1)': {
        'ET': {'mean': 1.27, 'cv': 44}, 'PT': {'mean': 1.27, 'cv': 44},
        'IT': {'mean': 1.27, 'cv': 44}, 'UT': {'mean': 1.27, 'cv': 44},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC22A7 (OAT2)': {
        'ET': {'mean': 1.25, 'cv': 75}, 'PT': {'mean': 1.25, 'cv': 75},
        'IT': {'mean': 1.25, 'cv': 75}, 'UT': {'mean': 1.25, 'cv': 75},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC22A9 (OAT7)': {
        'ET': {'mean': 1.45, 'cv': 51}, 'PT': {'mean': 1.45, 'cv': 51},
        'IT': {'mean': 1.45, 'cv': 51}, 'UT': {'mean': 1.45, 'cv': 51},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC47A1 (MATE1)': {
        'ET': {'mean': 0.146, 'cv': 51}, 'PT': {'mean': 0.146, 'cv': 51},
        'IT': {'mean': 0.146, 'cv': 51}, 'UT': {'mean': 0.146, 'cv': 51},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    },
    'SLC51A/B (OST-α/β)': {
        'ET': {'mean': 0.0, 'cv': 60}, 'PT': {'mean': 0.0, 'cv': 60},
        'IT': {'mean': 0.0, 'cv': 60}, 'UT': {'mean': 0.0, 'cv': 60},
        'turnover_mean': 1.00E-06, 'turnover_cv': 30
    }
}

def clean_parameter_name(name):
    """清理参数名称，移除特殊字符"""
    return re.sub(r'[^a-zA-Z0-9]', '', str(name))


def clean_transporter_name(name):
    """清理转运体名称，保留括号但移除其他特殊字符"""
    # 保留括号、字母、数字，移除其他特殊字符
    return re.sub(r'[^a-zA-Z0-9()]', '', str(name))

def sample_lognormal(mean, cv, size=1):
    """从对数正态分布采样"""
    if mean <= 0:
        return np.zeros(size)
    sigma_log = np.sqrt(np.log(1 + (cv/100)**2))
    mu_log = np.log(mean) - 0.5 * sigma_log**2
    return np.exp(np.random.normal(loc=mu_log, scale=sigma_log, size=size))

def generate_cyp_genotypes_and_phenotypes(pop_df, random_seed=42):
    """
    为每个人生成CYP酶的基因型和表型
    
    参数:
    pop_df: 包含id的人口统计学DataFrame
    random_seed: 随机种子
    
    返回:
    包含CYP基因型和表型的DataFrame
    """
    np.random.seed(random_seed)
    n_subjects = len(pop_df)
    results = {}
    results['id'] = pop_df['id'].values
    
    # 为每个CYP酶生成基因型和表型
    for cyp_name in CYP_FREQUENCIES.keys():
        frequencies = CYP_FREQUENCIES[cyp_name]
        abundances = CYP_ABUNDANCES[cyp_name]
        
        # 生成基因型
        phenotypes = list(frequencies.keys())
        phenotype_probs = list(frequencies.values())
        
        # 确保概率和为1
        total_prob = sum(phenotype_probs)
        if total_prob > 0:
            phenotype_probs = [p/total_prob for p in phenotype_probs]
        else:
            phenotype_probs = [1.0] + [0.0] * (len(phenotypes) - 1)
        
        # 为每个人分配基因型
        genotypes = np.random.choice(phenotypes, size=n_subjects, p=phenotype_probs)
        
        # 生成对应的表型数据
        abundances_values = []
        turnover_values = []
        
        for i, genotype in enumerate(genotypes):
            # 获取该基因型的丰度参数
            abundance_params = abundances[genotype]
            mean_abundance = abundance_params['mean']
            cv_abundance = abundance_params['cv']
            
            # 生成丰度值
            if mean_abundance > 0:
                abundance = sample_lognormal(mean_abundance, cv_abundance, 1)[0]
            else:
                abundance = 0.0
            
            abundances_values.append(abundance)
            
            # 生成周转时间（所有人使用相同的周转参数）
            if i == 0:  # 只在第一次计算周转时间
                turnover_mean = abundances['turnover_mean']
                turnover_cv = abundances['turnover_cv']
                turnover = sample_lognormal(turnover_mean, turnover_cv, 1)[0]
                turnover_values = [turnover] * n_subjects
        
        # 存储结果
        clean_cyp = clean_parameter_name(cyp_name)
        results[f'Liver_{clean_cyp}_genotype'] = genotypes
        results[f'Liver_{clean_cyp}_abundance_pmol_mg'] = abundances_values
        results[f'Liver_{clean_cyp}_turnover_1_h'] = turnover_values
    
    return pd.DataFrame(results)


def generate_transporter_genotypes_and_phenotypes(pop_df, random_seed=42):
    """
    生成转运体基因型和表型数据
    
    参数:
    pop_df: 包含id的人口统计学DataFrame
    random_seed: 随机种子
    
    返回:
    包含转运体基因型和表型的DataFrame
    """
    np.random.seed(random_seed)
    n_subjects = len(pop_df)
    results = {}
    results['id'] = pop_df['id'].values
    
    # 为每个转运体生成基因型和表型
    for transporter_name in TRANSPORTER_FREQUENCIES.keys():
        frequencies = TRANSPORTER_FREQUENCIES[transporter_name]
        abundances = TRANSPORTER_ABUNDANCES[transporter_name]
        
        # 生成基因型
        phenotypes = list(frequencies.keys())
        phenotype_probs = list(frequencies.values())
        
        # 确保概率和为1
        total_prob = sum(phenotype_probs)
        if total_prob > 0:
            phenotype_probs = [p/total_prob for p in phenotype_probs]
        else:
            phenotype_probs = [1.0] + [0.0] * (len(phenotypes) - 1)
        
        # 为每个人分配基因型
        genotypes = np.random.choice(phenotypes, size=n_subjects, p=phenotype_probs)
        
        # 生成对应的表型数据
        abundances_values = []
        turnover_values = []
        
        for i, genotype in enumerate(genotypes):
            # 获取该基因型的丰度参数
            abundance_params = abundances[genotype]
            mean_abundance = abundance_params['mean']
            cv_abundance = abundance_params['cv']
            
            # 生成丰度值
            if mean_abundance > 0:
                abundance = sample_lognormal(mean_abundance, cv_abundance, 1)[0]
            else:
                abundance = 0.0
            
            abundances_values.append(abundance)
            
            # 生成周转时间（所有人使用相同的周转参数）
            if i == 0:  # 只在第一次计算周转时间
                turnover_mean = abundances['turnover_mean']
                turnover_cv = abundances['turnover_cv']
                turnover = sample_lognormal(turnover_mean, turnover_cv, 1)[0]
                turnover_values = [turnover] * n_subjects
        
        # 存储结果
        clean_transporter = clean_transporter_name(transporter_name)
        results[f'Liver_{clean_transporter}_genotype'] = genotypes
        results[f'Liver_{clean_transporter}_abundance_pmol_mg'] = abundances_values
        results[f'Liver_{clean_transporter}_turnover_1_h'] = turnover_values
    
    return pd.DataFrame(results)


def generate_liver_enzymes_transporters_from_df(pop_df, random_seed=42):
    """
    生成肝脏酶和转运体参数的主函数
    
    参数:
    pop_df: 包含id的人口统计学DataFrame
    random_seed: 随机种子
    
    返回:
    包含肝脏酶和转运体参数的DataFrame
    """
    # 生成CYP酶数据
    cyp_df = generate_cyp_genotypes_and_phenotypes(pop_df, random_seed)
    
    # 生成转运体数据
    transporter_df = generate_transporter_genotypes_and_phenotypes(pop_df, random_seed)
    
    # 合并CYP酶和转运体数据
    result_df = pd.merge(cyp_df, transporter_df, on='id', how='left')
    
    return result_df

if __name__ == "__main__":
    # 测试代码
    test_df = pd.DataFrame({
        'id': [1, 2, 3, 4, 5],
        'age': [30, 40, 50, 60, 70],
        'sex': ['M', 'F', 'M', 'F', 'M']
    })
    
    result = generate_liver_enzymes_transporters_from_df(test_df)
    print("生成的CYP酶数据:")
    print(result.head())
    
    # 显示每个CYP酶的基因型分布
    for col in result.columns:
        if col.endswith('_genotype'):
            print(f"\n{col} 分布:")
            print(result[col].value_counts()) 